{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Knet CNN Example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "using Pkg; haskey(Pkg.installed(),\"Knet\") || Pkg.add(\"Knet\")\n",
    "using Knet\n",
    "True=true # so we can read the python params\n",
    "include(\"common/params.py\")\n",
    "gpu()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "OS: Linux\n",
      "Julia: 1.0.0\n",
      "Knet: 1.0.1+\n",
      "GPU: TITAN Xp, TITAN Xp, TITAN Xp, TITAN Xp, TITAN Xp, TITAN Xp, TITAN Xp, TITAN Xp, \n"
     ]
    }
   ],
   "source": [
    "println(\"OS: \", Sys.KERNEL)\n",
    "println(\"Julia: \", VERSION)\n",
    "println(\"Knet: \", Pkg.installed()[\"Knet\"])\n",
    "println(\"GPU: \", replace(read(`nvidia-smi --query-gpu=name --format=csv,noheader`,String),'\\n'=>\", \"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define model\n",
    "function initmodel(; atype=KnetArray, dtype=Float32, winit=xavier, binit=zeros)\n",
    "    w(dims...)=atype(winit(dtype,dims...))\n",
    "    b(dims...)=atype(binit(dtype,dims...))\n",
    "    return Any[\n",
    "        w(3,3,3,50), b(1,1,50,1),\n",
    "        w(3,3,50,50), b(1,1,50,1),\n",
    "        w(3,3,50,100), b(1,1,100,1),\n",
    "        w(3,3,100,100), b(1,1,100,1),\n",
    "        w(512,6400), b(512,1),\n",
    "        w(10,512), b(10,1)\n",
    "    ]\n",
    "end;"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define loss and its gradient\n",
    "function predict(w,x; pdrop=(0,0))\n",
    "    convbias(x,w,b) = conv4(w,x;padding=1) .+ b\n",
    "    fc(x,w,b) = w * mat(x) .+ b;\n",
    "    x = relu.(convbias(x,w[1],w[2]))\n",
    "    x = relu.(pool(convbias(x,w[3],w[4])))\n",
    "    x = dropout(x,pdrop[1])\n",
    "    x = relu.(convbias(x,w[5],w[6]))\n",
    "    x = relu.(pool(convbias(x,w[7],w[8])))\n",
    "    x = dropout(x,pdrop[1])\n",
    "    x = relu.(fc(x,w[9],w[10]))\n",
    "    x = dropout(x,pdrop[2])\n",
    "    return fc(x,w[11],w[12])\n",
    "end\n",
    "\n",
    "loss(w,x,y;o...)=nll(predict(w,x;o...),y) # nll: negative log likelihood\n",
    "lossgradient = grad(loss);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  0.070144 seconds (67.62 k allocations: 3.424 MiB)\n",
      "32×32×3×50000 Array{Float32,4}\n",
      "50000-element Array{UInt8,1}\n",
      "32×32×3×10000 Array{Float32,4}\n",
      "10000-element Array{UInt8,1}\n"
     ]
    }
   ],
   "source": [
    "# load data\n",
    "include(Knet.dir(\"data\",\"cifar.jl\"))\n",
    "@time (xtrn,ytrn,xtst,ytst,lbls)=cifar10()\n",
    "dtrn = minibatch(xtrn,ytrn,BATCHSIZE,shuffle=true,xtype=KnetArray)\n",
    "dtst = minibatch(xtst,ytst,BATCHSIZE,shuffle=false,xtype=KnetArray)\n",
    "for d in (xtrn,ytrn,xtst,ytst); println(summary(d)); end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "# prepare for training\n",
    "model = optim = nothing; Knet.gc() # Clear memory from last run\n",
    "model = initmodel()\n",
    "optim = optimizers(model, Momentum; lr=LR, gamma=MOMENTUM);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  6.474056 seconds (2.00 M allocations: 670.819 MiB, 1.59% gc time)\n"
     ]
    }
   ],
   "source": [
    "# cold start\n",
    "@time for (x,y) in dtrn\n",
    "    grads = lossgradient(model, x, y; pdrop=(0.25,0.5))\n",
    "    update!(model, grads, optim)\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "# prepare for training\n",
    "model = optim = nothing; Knet.gc() # Clear memory from last run\n",
    "model = initmodel()\n",
    "optim = optimizers(model, Momentum; lr=LR, gamma=MOMENTUM);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "┌ Info: Training...\n",
      "└ @ Main In[36]:2\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  5.831162 seconds (1.88 M allocations: 664.773 MiB, 1.54% gc time)\n",
      "  4.364551 seconds (1.88 M allocations: 664.628 MiB, 2.05% gc time)\n",
      "  4.277376 seconds (1.88 M allocations: 664.628 MiB, 2.06% gc time)\n",
      "  4.452553 seconds (1.88 M allocations: 664.628 MiB, 1.99% gc time)\n",
      "  4.311131 seconds (1.88 M allocations: 664.628 MiB, 1.93% gc time)\n",
      "  4.310649 seconds (1.88 M allocations: 664.628 MiB, 2.09% gc time)\n",
      "  4.327310 seconds (1.88 M allocations: 664.628 MiB, 1.92% gc time)\n",
      "  4.363336 seconds (1.88 M allocations: 664.628 MiB, 2.22% gc time)\n",
      "  4.308050 seconds (1.88 M allocations: 664.626 MiB, 1.98% gc time)\n",
      "  4.379973 seconds (1.88 M allocations: 664.626 MiB, 1.94% gc time)\n",
      " 44.936848 seconds (18.84 M allocations: 6.491 GiB, 1.96% gc time)\n"
     ]
    }
   ],
   "source": [
    "# 159s\n",
    "@info(\"Training...\")\n",
    "@time for epoch in 1:EPOCHS\n",
    "    @time for (x,y) in dtrn\n",
    "        grads = lossgradient(model, x, y; pdrop=(0.25,0.5))\n",
    "        update!(model, grads, optim)\n",
    "    end\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  0.551547 seconds (131.39 k allocations: 124.437 MiB, 4.39% gc time)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "0.7834535256410257"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# test accuracy 77.54\n",
    "@time accuracy(model,dtst,predict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Julia 1.0.0",
   "language": "julia",
   "name": "julia-1.0"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.0.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
